# -*- coding:utf-8 -*-
import datetime
import pandas as pd
from imblearn.pipeline import make_pipeline
from lightgbm import LGBMClassifier
from gwlsa_settings import net_params
import gwlsa_settings as GS
from utils import general_utils
from utils.landslide_utils import load_df_fromfile
from utils.general_utils import timer

def prepare_data_lightgbm():
    # 检查设置保存结果的各个文件夹是否存在，不存在则创建
    general_utils.check_data_settings_folder()

    train_df, val_df, test_df = load_df_fromfile(net_params['data_load_dir'], net_params['max_distance'],
                                                 net_params['resolution'], file_format=net_params['data_load_format'])

    train_val_df = pd.concat([train_df, val_df], ignore_index=False)
    all_df = pd.concat([train_val_df, test_df], ignore_index=False)

    X_all, y_all = all_df[net_params['x_column_names']], all_df[net_params['y_column_name']]
    GeoID = all_df[net_params['id_column']]

    # 合并的train_df和val_df作为训练数据集
    X_train, y_train = train_val_df[net_params['x_column_names']], train_val_df[net_params['y_column_name']]
    X_test, y_test = test_df[net_params['x_column_names']], test_df[net_params['y_column_name']]

    return X_train, y_train, X_test, y_test, GeoID, X_all, y_all

@timer
def exec_lightgbm(max_class_weight=2):
    X_train, y_train, X_test, y_test, GeoID, X_all, y_all = prepare_data_lightgbm()
    ###############################################################################
    # Classification using lightGBM classifier with and without sampling
    ###############################################################################
    print(f'Weighted LightGBM classifier performance:')
    for i in range(1, max_class_weight+1):
        wlgb = LGBMClassifier(random_state=0, n_jobs=-1, class_weight={0: 1, 1: i}, verbose=-1) #在调试模型或希望了解训练过程（如早期停止是否触发）时，建议将 verbose 设置回一个非负值（例如 verbose=1），以便获取有价值的信息。
        pipeline_wlgb = make_pipeline(wlgb)
        pipeline_wlgb.fit(X_train, y_train)
        y_pred_wlgb = pipeline_wlgb.predict(X_test)
        y_pred_wlgb_prob = pipeline_wlgb.predict_proba(X_test)[:, 1]
        fpr_wlgb, tpr_wlgb, auc_wlgb = general_utils.print_performance(y_test, y_pred_wlgb, y_pred_wlgb_prob, prefix_info=f'[class_weight 1: {i}]')

        # 保存整个数据集预测结果【包含train,val和test】到csv文件
        saved_dir = GS.PREDICTED_Y_DIR
        y_pred_wlgb = pipeline_wlgb.predict(X_all)
        y_pred_wlgb_prob = pipeline_wlgb.predict_proba(X_all)[:, 1]
        # 预测结果并保存到文件【测试数据集】
        weights_str = f'w1to{i}'
        t_name = datetime.datetime.today().strftime("%Y%m%d%H%M%S")
        save_csv_filename = f'{saved_dir}/LightGBM_{weights_str}_{t_name}.csv'
        general_utils.save_results(GeoID, y_pred_wlgb, y_pred_wlgb_prob, save_csv_filename, print_info=False)

        print('Done!')

    # cm_wlgb = confusion_matrix(y_test, y_pred_wlgb)
    # print(classification_report(y_test, y_pred_wlgb, target_names=['no-landslide', 'landslide']))
    # plot_all_confusion_matrix(cm_lgb, cm_wlgb)
    # plot_roc(fpr_lgb, tpr_lgb, auc_lgb, fpr_wlgb, tpr_wlgb, auc_wlgb)

if __name__ == '__main__':
    exec_lightgbm(max_class_weight=2)

